import time
import subprocess   # 用于开启多个进程
import psutil       # 使用psutil进行资源监控，主要获取cpu与内存占用情况
import time
from keras import backend as K
from keras.models import load_model
from tensorflow.python.saved_model import builder as saved_model_builder
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.saved_model.signature_def_utils_impl import predict_signature_def
##################################################################################################################
# 多模型多进程的训练逻辑:
# 开启第一个模型训练进程，进入训练状态后开始检测占用资源是否小于CPU/内存正常负载值。
# 小于CPU/内存正常负载值，则开启第二个模型训练任务。否则，开始检测占用资源是否大于CPU/内存危险负载值，
# 如果大于，则kill掉这个进程，否则，说明占用率处在正常负载值与危险负载值之间，此时，等待该模型训练进程结束，再自动开启下一个模型训练进程。
##################################################################################################################

# 设定CPU与内存的正常和危险占用阈值
CPU_NOR_LIMIT = MEM_NOR_LIMIT = 55
CPU_DAN_LIMIT = MEM_DAN_LIMIT = 95

# 模型训练脚本列表
model_train_list = ["python movie_model_train.py", "python beauty_model_train.py"]

# 创建subp的列表容器，用于装载子进程
subp = []

def detect_cpu_mem():
    """检测CPU和内存的占用率"""
    print("进行mem和cpu检测:")
    # 内存检测
    mem = psutil.virtual_memory().percent   # 内存信息的获取
    # psutil检测cpu时间隔至少3s以上
    cpu = psutil.cpu_percent(3)             # 获取cpu的使用率
    print("当前内存占用率:" + str(mem) + "%")
    print("当前CPU占用率:" + str(cpu) + "%")
    return mem, cpu


def single_model_train(model):
    """开启单个模型的训练"""
    p = subprocess.Popen(model, shell=True)
    # 等待3秒预估模型进入训练状态，即资源占用趋于稳定。
    time.sleep(3)
    # 进行资源检测
    mem, cpu = detect_cpu_mem()

    # 内存和CPU同时小于正常负载值，则任其继续运行，并装入列表
    if mem < MEM_NOR_LIMIT and cpu < CPU_NOR_LIMIT:
        subp.append(p)
        print("该模型进入正常训练过程，并可以开启下一模型训练！")
    else:
        # 判断是否大于危险负载值，若大于，将kill该进程，
        # 否则等待该进程结束，再进行其他训练任务。
        if mem > MEM_DAN_LIMIT or cpu > CPU_DAN_LIMIT:
            p.kill()
            print("该模型没有进入正常训练过程！")
        else:
            p.wait()
            print("该模型进入正常训练过程, 但不要开启下一个模型训练!")


def start_multiprocess_train():
    """开启多进程训练"""
    print("启动多模型训练：")

    # 遍历启动模型的命令，准备循环开启训练进程
    for i, model in enumerate(model_train_list):
        print("__________________________")
        print("正在启动第" + str(i+1) + "个模型：")
        # 启动模型训练
        single_model_train(model)
    else:
        # 所有装入列表的进程都会等待其自然结束后才会停止该函数所在的进程
        print("正在等待所有模型训练结束!")
        list(map(lambda x: x.wait(), subp))
        print("完成!")



######################################################################################################
                     #多模型多线程预测:
                     #学习并实现了多模型多线程预测的三步曲:
                     #第一步: 将h5格式的模型转化成pb格式。
                     #第二步: 使用docker启动tensorflow-serving微服务。
#                    第三步: 多线程调用微服务并处理结果。
######################################################################################################

def to_savedmodel(h5_model_path, pb_model_path):
  """将h5模型转化成tensorflow的pb格式模型"""

  model = load_model(h5_model_path)
  builder = saved_model_builder.SavedModelBuilder(pb_model_path)

  signature = predict_signature_def(
      inputs={'input': model.inputs[0]}, outputs={'income': model.outputs[0]})

  with K.get_session() as sess:
    builder.add_meta_graph_and_variables(
        sess=sess,
        tags=[tag_constants.SERVING],
        signature_def_map={
            signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature
        })
    builder.save()